import comparisonTest.GetUnsafeUtil;
import sun.misc.Unsafe;

import java.lang.reflect.Field;
import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.*;
import java.util.concurrent.locks.ReentrantLock;

public class LockFreeHashMap<K, V> implements ConcurrentMap<K, V> {
    private static final int DEFAULT_CAPACITY = 1024;
    private static final boolean PRINT_CAS_FAIL = false;
    private static final boolean LOCK_ON_CAS_FAIL = false;

    static class Node<K, V> implements Map.Entry<K, V> {
        private static final Unsafe unsafe;
        private static final long nextFieldOffset;
        private static final long valueFieldOffset;

        final int hash;
        final K key;
        volatile V value;
        volatile Node<K, V> next;

        static {
            try {
                unsafe = GetUnsafeUtil.getUnsafe();
                Field nextField = Node.class.getDeclaredField("next");
                nextFieldOffset = unsafe.objectFieldOffset(nextField);
                Field valueField = Node.class.getDeclaredField("value");
                valueFieldOffset = unsafe.objectFieldOffset(valueField);
            } catch (Exception ex) {
                throw new Error(ex);
            }
        }

        public Node(K key, V value, int hash, Node<K, V> next) {
            this.key = key;
            this.value = value;
            this.hash = hash;
            this.next = next;
        }

        @Override
        public K getKey() {
            return key;
        }

        @Override
        public V getValue() {
            return value;
        }

        @Override
        public V setValue(V newValue) {
            V old = value;
            value = newValue;
            return old;
        }

        public boolean compareAndSwapNext(Node<K, V> expect, Node<K, V> update) {
            return unsafe.compareAndSwapObject(this, nextFieldOffset, expect, update);
        }

        private boolean compareAndSwapValue(V expect, V update) {
            return unsafe.compareAndSwapObject(this, valueFieldOffset, expect, update);
        }

        public boolean reset(V value) {
            if (!isRemoved()) return false;
            return compareAndSwapValue(null, value);
        }

        public boolean setRemoved() {
            // if (this.removed.compareAndSet(false, true)) { ... }  //有没有可能value的初始值会对cas有影响？感觉应该没有
            return compareAndSwapValue(value, null);
        }

        public boolean isRemoved() {
            return this.value == null;
        }

        public boolean equalsKey(Object key, int hash, boolean allowRemoved) {
            return this.hash == hash &&
                    (allowRemoved || !isRemoved()) &&
                    (key == this.key || key.equals(this.key));
        }

        public boolean equalsKey(Object key, int hash) {
            return this.hash == hash &&
                    this.value != null &&
                    (key == this.key
                            || key.equals(this.key));
        }

        public void setNext(Node<K, V> next) {
            this.next = next;
        }
    }

    private final AtomicReferenceArray<Node<K, V>> table;
    private final AtomicIntegerArray counterCells;
    private final ReentrantLock[] failLocks;
    private final int capacity;
    private final int hashMask;

    public LockFreeHashMap() {
        this(DEFAULT_CAPACITY);
    }

    public LockFreeHashMap(int capacity) {
        this.capacity = capacity;
        this.hashMask = capacity - 1;
        this.table = new AtomicReferenceArray<>(capacity);
        this.counterCells = new AtomicIntegerArray(capacity);
        this.failLocks = new ReentrantLock[capacity];
        for (int i = 0; i < capacity; i++) {
            this.failLocks[i] = new ReentrantLock();
        }
    }

    @Override
    public V get(Object key) {
        if (key == null) return null;
        int hash = getHash(key.hashCode());
        int index = hash & hashMask;
        Node<K, V> node = table.get(index);
        while (node != null && !node.equalsKey(key, hash)) {
            node = node.next;
        }
        if (node == null) return null;
        else return node.getValue();
    }

    @Override
    public V put(K key, V value) {
        if (key == null || value == null)
            throw new IllegalArgumentException("Key or value cannot be null");
        int hash = getHash(key.hashCode());
        int index = hash & hashMask;
        Node<K, V> newNode = null;
        int step = 1;
        int currentCount = counterCells.get(index);
        if (currentCount > 1)
            step = ThreadLocalRandom.current().nextInt(currentCount);
        boolean locked = false;
        if (LOCK_ON_CAS_FAIL && failLocks[index].isLocked()) {
            failLocks[index].lock();
            locked = true;
        }
        try {
            Node<K, V> node = table.get(index);
            if (node == null) {
                newNode = new Node<>(key, value, hash, null);
                while (true) {
                    if (table.compareAndSet(index, null, newNode)) {
                        incrementCount(index);
                        return null;
                    } else if ((node = table.get(index)) != null) {
                        break;
                    }
                }
            }
            boolean eq = false;
            for (int i = 0; i < step; i++) {
                if (node.equalsKey(key, hash, true)) {
                    eq = true;
                    break;
                }
                Node<K, V> next = node.next;
                if (next == null) break;
                node = next;
            }
            if (!eq) {
                Node<K, V> node2 = node.next;
                while (node2 != null) {
                    if (node2.equalsKey(key, hash, true)) {
                        node = node2;
                        eq = true;
                        break;
                    }
                    node2 = node2.next;
                }
            }
            if (eq) {
                if (node.isRemoved()) {
                    if (node.reset(value))
                        incrementCount(index);
                    return null;
                } else {
                    return node.setValue(value);
                }
            }
            if (newNode == null)
                newNode = new Node<>(key, value, hash, null);
            while (true) {
                Node<K, V> next = node.next;
                newNode.setNext(next);
                if (node.compareAndSwapNext(next, newNode)) {
                    break;
                } else {
                    if (LOCK_ON_CAS_FAIL) {
                        failLocks[index].lock();
                        locked = true;
                    }
                    if (PRINT_CAS_FAIL)
                        System.out.println("CAS failed for " + newNode.getKey());
                }
            }
            // 备注：这种方式有可能会对同一个key插入了多个node，删除时需要注意
            // 但get不受影响，因为会优先取靠近头节点的
            incrementCount(index);
            return null;
        } finally {
            if (locked) {
                failLocks[index].unlock();
            }
        }
    }

    @Override
    public V remove(Object key) {
        if (key == null) return null;
        V oldValue = null;
        int hash = getHash(key.hashCode());
        int index = hash & hashMask;
        int cnt = 0;
        Node<K, V> node = table.get(index);
        while (node != null) {
            if (node.equalsKey(key, hash)) {
                if (oldValue == null)
                    oldValue = node.getValue();
                if (node.setRemoved())
                    cnt++;
            }
            node = node.next;
        }
        if (cnt > 0)
            counterCells.getAndAdd(index, -cnt);
        return oldValue;
    }

    private void incrementCount(int index) {
        counterCells.getAndIncrement(index);
    }

    public void truncate() {
        int cnt = 0;
        for (int i = 0; i < table.length(); i++) {
            failLocks[i].lock();
            Node<K, V> node = table.get(i);
            while (node != null) {
                Node<K, V> next = node.next;
                if (next == null) break;
                if (next.isRemoved()) {
                    Node<K, V> next_next = next.next;
                    node.next = next_next;
                    cnt++;
                }
                node = node.next;
            }
            failLocks[i].unlock();
        }
        System.out.println("Truncate removed " + cnt + " nodes");
    }

    @Override
    public void putAll(Map<? extends K, ? extends V> m) {

    }

    @Override
    public void clear() {

    }

    @Override
    public int size() {
        int size = 0;
        for (int i = 0; i < counterCells.length(); i++) {
            size += counterCells.get(i);
        }
        return Math.max(size, 0);
    }

    @Override
    public boolean isEmpty() {
        for (int i = 0; i < counterCells.length(); i++) {
            if (counterCells.get(i) > 0) return false;
        }
        return true;
    }

    @Override
    public boolean containsKey(Object key) {
        return false;
    }

    @Override
    public boolean containsValue(Object value) {
        return false;
    }


    @Override
    public Set<K> keySet() {
        return null;
    }

    @Override
    public Collection<V> values() {
        return null;
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        return null;
    }

    @Override
    public V putIfAbsent(K key, V value) {
        return null;
    }

    @Override
    public boolean remove(Object key, Object value) {
        return false;
    }

    @Override
    public boolean replace(K key, V oldValue, V newValue) {
        return false;
    }

    @Override
    public V replace(K key, V value) {
        return null;
    }

    static int getHash(int hashCode) {
        return (hashCode ^ (hashCode >>> 16)) & 0x7fffffff;
    }
}
